import serial
import minimalmodbus
import time
import threading
import sys

class LingkongMotor:
    # 保持原有代码不变
    def __init__(self, serial_port, motor_id=0x01, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.ser = serial_port
        self.motor_id = motor_id
        self.low_speed = 20
        self.zero_angle_target = 0
        self.lock = lock or threading.Lock()

    def _calc_checksum(self, data):
        return sum(data) & 0xFF

    def _send_command(self, cmd, data_bytes=b'', read_len=32, post_delay=0.002):
        frame_cmd = [0x3E, cmd, self.motor_id, len(data_bytes)]
        frame_cmd_checksum = self._calc_checksum(frame_cmd)
        frame_cmd.append(frame_cmd_checksum)

        full_frame = bytes(frame_cmd) + data_bytes
        if data_bytes:
            data_checksum = self._calc_checksum(data_bytes)
            full_frame += bytes([data_checksum])

        with self.lock:
            try:
                try:
                    self.ser.reset_input_buffer()
                    self.ser.reset_output_buffer()
                except Exception:
                    pass

                self.ser.write(full_frame)
                self.ser.flush()
                time.sleep(post_delay)
                resp = self.ser.read(read_len)
                return resp
            except Exception:
                raise

    def get_current_angle(self):
        try:
            response = self._send_command(0x92, b'', read_len=32)
        except Exception as e:
            print(f"[Motor] _send_command 出错: {e}")
            return None

        if response and len(response) >= 14:
            try:
                angle_bytes = response[5:13]
                current_angle_raw = int.from_bytes(angle_bytes, byteorder='little', signed=True)
                return current_angle_raw * 0.01
            except Exception as e:
                print(f"[Motor] 解析角度失败: {e}")
                return None
        return None

    def move_to_zero(self, timeout=300):
        print("开始回归零位...")
        target_bytes = int(self.zero_angle_target).to_bytes(8, byteorder='little', signed=True)
        speed_value = int(self.low_speed * 100)
        speed_bytes = int(speed_value).to_bytes(4, byteorder='little', signed=False)

        try:
            self._send_command(0xA4, target_bytes + speed_bytes)
        except Exception as e:
            print(f"[Motor] 发送回零命令失败: {e}")
            return False

        start_time = time.time()
        while time.time() - start_time < timeout:
            current = self.get_current_angle()
            if current is not None and abs(current) <= 0.1:
                print("已到达零位")
                return True
            time.sleep(0.2)
        print("回零超时")
        return False

    def rotate_360(self, timeout=300):
        print("开始旋转360度...")
        target = 36000
        target_bytes = int(target).to_bytes(8, byteorder='little', signed=True)
        speed_value = int(self.low_speed * 10)
        speed_bytes = int(speed_value).to_bytes(4, byteorder='little', signed=False)

        try:
            self._send_command(0xA4, target_bytes + speed_bytes)
        except Exception as e:
            print(f"[Motor] 发送旋转命令失败: {e}")
            return False

        start_time = time.time()
        while time.time() - start_time < timeout:
            current = self.get_current_angle()
            if current is not None:
                ang = current % 360
                if abs(ang - 0) <= 0.1:
                    print("已完成360度旋转（归一化为0）")
                    return True
            time.sleep(0.2)
        print("旋转超时")
        return False

    def close(self):
        pass


class LE100485Sensor:
    # 保持原有代码不变
    def __init__(self, serial_port, sensor_address=2, lock=None):
        if serial_port is None:
            raise ValueError("必须提供共享 serial_port 实例")
        self.serial_port = serial_port
        self.lock = lock or threading.Lock()
        self.sensor_address = sensor_address

        self.instrument = minimalmodbus.Instrument(serial_port, slaveaddress=self.sensor_address)
        self.instrument.mode = minimalmodbus.MODE_RTU
        self.instrument.serial = self.serial_port
        self.instrument.close_port_after_each_call = False
        self.instrument.clear_buffers_before_each_transaction = True
        self.instrument.debug = False

        self.distance_register = 0x003B
        self.resolution = 0.001

    def read_distance(self, register_address=None):
        if register_address is None:
            register_address = self.distance_register
        with self.lock:
            try:
                raw_data = self.instrument.read_register(
                    registeraddress=register_address,
                    functioncode=3,
                    signed=False
                )
                return raw_data * self.resolution
            except Exception as e:
                print(f"[Sensor] 读取失败: {e}")
                return None

    def close(self):
        pass


class StepperMotor:
    """步进电机控制类（地址3）"""
    def __init__(self, serial_port, motor_id=3, lock=None,细分=10000):
        self.ser = serial_port
        self.motor_id = motor_id  # 步进电机地址3
        self.细分 = 细分  # 细分参数，默认为10000
        self.lock = lock or threading.Lock()

    def _send_modbus_command(self, register, value):
        """发送Modbus RTU写单个寄存器命令（功能码06）"""
        with self.lock:
            try:
                # 构建命令帧：地址 + 功能码06 + 寄存器高8位 + 寄存器低8位 + 数据高8位 + 数据低8位 + 校验和
                cmd = bytearray()
                cmd.append(self.motor_id)  # 设备地址
                cmd.append(0x06)  # 功能码06（写单个寄存器）
                
                # 寄存器地址（2字节）
                cmd.extend(register.to_bytes(2, byteorder='big'))
                
                # 写入值（2字节）
                cmd.extend(value.to_bytes(2, byteorder='big'))
                
                # 计算CRC校验和
                crc = self._calc_crc(cmd)
                cmd.extend(crc)
                
                # 发送命令
                self.ser.reset_input_buffer()
                self.ser.reset_output_buffer()
                self.ser.write(cmd)
                self.ser.flush()
                
                # 等待响应（8字节）
                time.sleep(0.01)
                resp = self.ser.read(8)
                return len(resp) == 8  # 响应正确返回True
            except Exception as e:
                print(f"[Stepper] 发送命令失败: {e}")
                return False

    def _calc_crc(self, data):
        """计算Modbus CRC16校验和"""
        crc = 0xFFFF
        for byte in data:
            crc ^= byte
            for _ in range(8):
                if crc & 0x0001:
                    crc >>= 1
                    crc ^= 0xA001
                else:
                    crc >>= 1
        return crc.to_bytes(2, byteorder='little')

    def rotate_cycles(self, cycles, speed=300, accel=50, decel=50, timeout=300):
        """
        控制步进电机旋转指定圈数
        cycles: 旋转圈数（正数为正转，负数为反转）
        """
        print(f"开始控制步进电机旋转 {cycles} 圈...")
        
        # 计算目标位置（1圈 = 细分值）
        target_position = int(cycles * self.细分)
        # 设置电流（0191）
        if not self._send_modbus_command(0x0191, 0x0028):
            print("设置PR0电流失败")
            return False
        
        # 设置PR0为绝对位置模式（0001）
        if not self._send_modbus_command(0x6200, 0x0001):
            print("设置PR0模式失败")
            return False
        
        # 处理目标位置（分高低位）
        position_high = (target_position >> 16) & 0xFFFF
        position_low = target_position & 0xFFFF
        if not self._send_modbus_command(0x6201, position_high):
            print("设置PR0位置高位失败")
            return False
        if not self._send_modbus_command(0x6202, position_low):
            print("设置PR0位置低位失败")
            return False
        
        # 设置速度、加速度、减速度
        if not self._send_modbus_command(0x6203, speed):
            print("设置PR0速度失败")
            return False
        if not self._send_modbus_command(0x6204, accel):
            print("设置PR0加速度失败")
            return False
        if not self._send_modbus_command(0x6205, decel):
            print("设置PR0减速度失败")
            return False
        
        # 触发PR0运行
        if not self._send_modbus_command(0x6002, 0x0010):
            print("触发PR0运行失败")
            return False
        
        # 等待完成（这里简化处理，实际可根据电机反馈优化）
        start_time = time.time()
        while time.time() - start_time < timeout:
            time.sleep(1)  # 每1秒检查一次
            # 实际应用中可通过读取电机状态寄存器判断是否完成
            # 这里暂时用延时估算，根据圈数和速度调整
            estimated_time = abs(cycles * self.细分) / (speed * 1.0)  # 估算时间
            if time.time() - start_time > estimated_time + 2:  # 加2秒缓冲
                print(f"步进电机完成 {cycles} 圈旋转")
                return True
        
        print("步进电机旋转超时")
        return False

    def emergency_stop(self):
        """急停命令"""
        return self._send_modbus_command(0x6002, 0x0040)


class DataCollector:
    # 保持原有代码不变
    def __init__(self, motor, sensor):
        self.motor = motor
        self.sensor = sensor
        self.data = []
        self.running = False
        self.thread = None

    def _collect_loop(self, interval):
        while self.running:
            angle = self.motor.get_current_angle()
            distance = self.sensor.read_distance()

            if angle is not None and distance is not None:
                normalized_angle = angle % 360
                self.data.append((normalized_angle, distance))
                print(f"角度: {normalized_angle:.2f}°  距离: {distance:.5f}mm", end='\r', flush=True)

            time.sleep(interval)

    def start(self, interval=0.05):
        if self.running:
            return
        self.data = []
        self.running = True
        self.thread = threading.Thread(target=self._collect_loop, args=(interval,), daemon=True)
        self.thread.start()

    def stop(self):
        self.running = False
        if self.thread:
            self.thread.join(timeout=2)

    def save_data(self, filename="measurement_data.csv"):
        with open(filename, 'w', encoding='utf-8') as f:
            f.write("角度(度),距离(mm)\n")
            for angle, distance in self.data:
                f.write(f"{angle:.4f},{distance:.5f}\n")
        print(f"\n数据已保存到 {filename}，共 {len(self.data)} 条记录")


def main():
    PORT = 'COM6'
    BAUDRATE = 115200
    MOTOR_ID = 1  # 关节电机地址
    SENSOR_ADDRESS = 2  # 激光传感器地址
    STEPPER_ID = 3  # 步进电机地址
    STEPPER_CYCLES = 2  # 步进电机需要旋转的圈数（可修改）

    shared_serial = None
    collector = None

    try:
        # 初始化共享串口
        shared_serial = serial.Serial(
            port=PORT,
            baudrate=BAUDRATE,
            parity=serial.PARITY_NONE,
            stopbits=serial.STOPBITS_ONE,
            bytesize=serial.EIGHTBITS,
            timeout=1,
            write_timeout=5
        )

        # 配置RS485模式
        try:
            if hasattr(serial, "rs485"):
                rs485_settings = serial.rs485.RS485Settings(
                    rts_level_for_tx=True,
                    rts_level_for_rx=False,
                    delay_before_tx=0.002,
                    delay_before_rx=0.002
                )
                shared_serial.rs485_mode = rs485_settings
                print("[Main] 已启用 serial.rs485_mode")
        except Exception as e:
            print(f"[Main] 无法启用 rs485_mode（忽略）: {e}")

        serial_lock = threading.Lock()

        # 初始化设备
        motor = LingkongMotor(serial_port=shared_serial, motor_id=MOTOR_ID, lock=serial_lock)
        sensor = LE100485Sensor(serial_port=shared_serial, sensor_address=SENSOR_ADDRESS, lock=serial_lock)
        stepper = StepperMotor(serial_port=shared_serial, motor_id=STEPPER_ID, lock=serial_lock)
        collector = DataCollector(motor, sensor)

        # 第一步：控制步进电机旋转指定圈数
        if not stepper.rotate_cycles(STEPPER_CYCLES):
            print("步进电机旋转失败，程序终止")
            return

        # 第二步：关节电机回零
        if not motor.move_to_zero():
            print("关节电机无法回归零位，程序终止")
            return

        # 第三步：启动数据采集并控制关节电机旋转
        collector.start(interval=0.05)
        if not motor.rotate_360():
            print("关节电机旋转失败，停止采集")
            collector.stop()
            return

        # 第四步：停止采集并保存数据
        collector.stop()
        collector.save_data()

        # 第五步：关节电机回零
        motor.move_to_zero()

        print("测量流程完成")

    except serial.SerialException as e:
        print(f"串口通信错误: {e}")
    except Exception as e:
        print(f"发生错误: {e}")
    finally:
        if collector:
            collector.stop()
        if shared_serial and shared_serial.is_open:
            try:
                shared_serial.close()
                print("串口已关闭")
            except Exception as e:
                print(f"关闭串口出错: {e}")

if __name__ == "__main__":
    main()